#include <stdlib.h>
#include <stdio.h>
#include <math.h>
#include <assert.h>
#include <mpi.h>
#include <gsl/gsl_math.h>
#include <gsl/gsl_sort_float.h>
#include <gsl/gsl_heapsort.h>
#ifdef USEOPENMP
#include <omp.h>
#endif

#include "halofinder.h"

typedef struct {
  int index;
  float dens;
} Peak;

#ifdef USESPHSMOOTHPEAKFINDER
static double sphKernel(double r, double h)
{
  double x = r/h;
  
  if(x > 1.0)
    return 0.0;
  else if(x >= 0.5)
    return 2.0*(1.0 - x)*(1.0 - x)*(1.0 - x)*8.0/M_PI/h/h/h;
  else
    return (1.0 - 6.0*x*x + 6.0*x*x*x)*8.0/M_PI/h/h/h;
}

static float getpeakdens_sphKernel(float centPos[3], int i, float *nbrsRad2, int *nbrsInd, int NumNbrs, int periodic, kdTreeData *td)
{
  float meanDens;
  int retNumNbrs;
  double sphOverDens,sphSmoothingLength;
  int j,loopstart;
  
  meanDens = (float) (haloFinderData.OmegaM*RHO_CRIT);
  
  retNumNbrs = get_knnbrs_kdtree(centPos,periodic,haloFinderData.BoxLengthSnapshot,
				 nbrsRad2,nbrsInd,NumNbrs,PartsPx,PartsPy,PartsPz,td);
  
  //remove part from nbrs
  loopstart = 0;
  for(j=0;j<retNumNbrs;++j)
    if(nbrsInd[j] == i)
      {
	loopstart = nbrsInd[0];
	nbrsInd[0] = nbrsInd[j];
	nbrsInd[j] = loopstart;
                
	sphSmoothingLength = nbrsRad2[0];
	nbrsRad2[0] = nbrsRad2[j];
	nbrsRad2[j] = (float) (sphSmoothingLength);
                
	loopstart = 1;
                
	break;
      }
          
  if(loopstart == 0)
    fprintf(stderr,"%d: loopstart is zero - particle %d did not find itself as nearest neighbor\n",ThisTask,i);
            
  //get max distance
  sphSmoothingLength = nbrsRad2[loopstart];
  for(j=loopstart+1;j<retNumNbrs;++j)
    if(nbrsRad2[j] > sphSmoothingLength)
      sphSmoothingLength = nbrsRad2[j];
  sphSmoothingLength = sqrt(sphSmoothingLength);
  
  //get overdens
  sphOverDens = 0.0;
  for(j=loopstart;j<retNumNbrs;++j)
    sphOverDens += sphKernel(sqrt(nbrsRad2[j]),sphSmoothingLength);
  sphOverDens = sphOverDens*haloFinderData.PartMass/meanDens;
  
  return (float) (sphOverDens);
}
#else
static float getpeakdens_tophat(float centPos[3], float **nbrsRad2, int **nbrsInd, int *maxNumNbrs, int periodic, kdTreeData *td)
{
  float searchRadius,searchVolume,meanDens;
  int NumNbrs;
  
  meanDens = (float) (haloFinderData.OmegaM*RHO_CRIT);
  searchRadius = (float) (3.0*haloFinderData.SofteningLength);
  searchVolume = (float) (4.0/3.0*M_PI*pow(searchRadius,3.0));
  
  NumNbrs = get_nnbrs_kdtree(centPos,searchRadius,periodic,haloFinderData.BoxLengthSnapshot,
			     nbrsRad2,nbrsInd,maxNumNbrs,PartsPx,PartsPy,PartsPz,td);
  
  return (float) (((double) NumNbrs)*haloFinderData.PartMass/searchVolume/meanDens);
}
#endif

void peakfinder(void)
{
  int i,peakInd,j;
  int periodic,maxNumNbrs;
  float peakDens;
  Halo *tmpHalos;
  float centPos[3];
  float *nbrsRad2;
  int *nbrsInd;
  Peak *Peaks;
  int NumPeaks,NumToTest,NumTested;
  int totNumPeaksFound;
  
  //setup
#ifdef LIGHTCONE
  periodic = 0;
#else
  periodic = 1;
#endif
  
  //look at all particles and find ones with highest overdens in radius of 3 softening lengths
#ifdef DEBUG
  if(ThisTask == 0)
    {
#ifdef USESPHSMOOTHPEAKFINDER
      fprintf(stderr,"finding peaks with SPH smoothing kernal\n");
#else
      fprintf(stderr,"finding peaks with top-hat smoothing\n");
#endif
    }
#endif
  
  NumHalos = 0;
  Halos = NULL;

#ifdef USEOPENMP    
#pragma omp parallel default(none) \
  shared(NumPartsInDomain,NumParts,PartsPx,PartsPy,PartsPz,PartsInDomainMark,periodic,haloFinderData,PartskdTreeData,ThisTask,stderr,NumHalos,Halos) \
  private(NumPeaks,Peaks,maxNumNbrs,peakDens,nbrsRad2,nbrsInd,i,centPos,peakInd,j,tmpHalos,NumToTest,NumTested,totNumPeaksFound)
#endif
  {
    //init halos
#ifdef USEOPENMP
    NumPeaks = (int) (IO_CHUNK_SIZE_MB*1024.0*1024.0/((double) (omp_get_num_threads()))/((double) (sizeof(Peak))));
#else
    NumPeaks = (int) (IO_CHUNK_SIZE_MB*1024.0*1024.0/((double) (sizeof(Peak))));
#endif
    Peaks = (Peak*)malloc(sizeof(Peak)*NumPeaks);
    assert(Peaks != NULL);
    
    //set max num nbrs
#ifndef USESPHSMOOTHPEAKFINDER
    maxNumNbrs = 100000;
#else
    maxNumNbrs = 65;
#endif
    nbrsRad2 = (float*)malloc(sizeof(float)*maxNumNbrs);
    assert(nbrsRad2 != NULL);
    nbrsInd = (int*)malloc(sizeof(int)*maxNumNbrs);
    assert(nbrsInd != NULL);
    
    peakInd = 0;
#ifdef USEOPENMP
    NumToTest = NumPartsInDomain/omp_get_num_threads();
#else
    NumToTest = NumPartsInDomain;
#endif
    NumTested = 0;

    totNumPeaksFound = 0;
    
#ifdef USEOPENMP
#pragma omp for schedule(guided)
#endif
    for(i=0;i<NumParts;++i)
      {

#ifdef DEBUG
	if(NumTested%(NumToTest/50) == 0 && ThisTask == 0 && NumTested > 0)
          fprintf(stderr,"%d: %d of %d (%.2lf percent) done in peak finding, %d peaks found, %.2lf percent of partciles\n",
		  ThisTask,NumTested,NumToTest,((double) NumTested)/((double) NumToTest)*100.0,totNumPeaksFound,((double) totNumPeaksFound)/((double) NumTested)*100.0);
#endif
	
	if(PartsInDomainMark[i])
	  {
	    ++NumTested;
	    
	    centPos[0] = PartsPx[i];
	    centPos[1] = PartsPy[i];
	    centPos[2] = PartsPz[i];
	    
#ifndef USESPHSMOOTHPEAKFINDER
	    peakDens = getpeakdens_tophat(centPos,&nbrsRad2,&nbrsInd,&maxNumNbrs,periodic,PartskdTreeData);
#else
	    peakDens = getpeakdens_sphKernel(centPos,i,nbrsRad2,nbrsInd,maxNumNbrs,periodic,PartskdTreeData);
#endif
	    
	    if(peakDens > haloFinderData.MinimumPeakOverdensWRTMeanDens)
	      {
		Peaks[peakInd].index = i;
		Peaks[peakInd].dens = peakDens;
		
		++totNumPeaksFound;
		++peakInd;
	
		if(peakInd == NumPeaks)
		  {	
#ifdef USEOPENMP
#pragma omp critical
#endif
		    {
		      tmpHalos = (Halo*)realloc(Halos,sizeof(Halo)*(NumHalos + NumPeaks));
		      assert(tmpHalos != NULL);
		      Halos = tmpHalos;
		      
		      for(j=0;j<NumPeaks;++j)
			{
			  Halos[NumHalos+j].index = Peaks[j].index;
			  Halos[NumHalos+j].mass = Peaks[j].dens;
			  Halos[NumHalos+j].radius = 0.0;
			  Halos[NumHalos+j].vcirc = 0.0;
			  Halos[NumHalos+j].pos[0] = PartsPx[Peaks[j].index];
			  Halos[NumHalos+j].pos[1] = PartsPy[Peaks[j].index];
			  Halos[NumHalos+j].pos[2] = PartsPz[Peaks[j].index];
			  Halos[NumHalos+j].vel[0] = PartsPx[Peaks[j].index];
			  Halos[NumHalos+j].vel[1] = PartsPy[Peaks[j].index];
			  Halos[NumHalos+j].vel[2] = PartsPz[Peaks[j].index];
			}
		      
		      NumHalos += NumPeaks;
		      peakInd = 0;
		    } /* end of critical region */ 
		  } 
	      }
	  }
      }

    if(peakInd > 0)
      {
#ifdef USEOPENMP
#pragma omp critical
#endif
	{    
	  tmpHalos = (Halo*)realloc(Halos,sizeof(Halo)*(NumHalos + peakInd));
	  assert(tmpHalos != NULL);
	  Halos = tmpHalos;
	  
	  for(j=0;j<peakInd;++j)
	    {
	      Halos[NumHalos+j].index = Peaks[j].index;
	      Halos[NumHalos+j].mass = Peaks[j].dens;
	      Halos[NumHalos+j].radius = 0.0;
	      Halos[NumHalos+j].vcirc = 0.0;
	      Halos[NumHalos+j].pos[0] = PartsPx[Peaks[j].index];
	      Halos[NumHalos+j].pos[1] = PartsPy[Peaks[j].index];
	      Halos[NumHalos+j].pos[2] = PartsPz[Peaks[j].index];
	      Halos[NumHalos+j].vel[0] = PartsPx[Peaks[j].index];
	      Halos[NumHalos+j].vel[1] = PartsPy[Peaks[j].index];
	      Halos[NumHalos+j].vel[2] = PartsPz[Peaks[j].index];
	    }
	  
	  NumHalos += peakInd;
	  peakInd = 0;
	} /* end of critical region */
      } 
    
    free(Peaks);
    free(nbrsRad2);
    free(nbrsInd);
  } /* end of parallel region */
  
#ifdef DEBUG
  fprintf(stderr,"%d: found %ld peaks\n",ThisTask,NumHalos);
#endif
}

void peakpercolation(void)
{
  int i,j,periodic,NumNbrs,maxNumNbrs,haloInd;
  kdTreeData *td;
  float searchRadius;
  Halo *tmpHalos;
  short int *keep;
  float *px,*py,*pz;
  float *nbrsRad2;
  int *nbrsInd;
  
#ifdef DEBUG
  if(ThisTask == 0)
    fprintf(stderr,"doing peak percolation\n");
#endif
  
  //setup - sort halos by mass, make them into parts for nnbrs finding, use keep vector to mark halos to keep during percolation
  sort_halos_mass();
    
  px = (float*)malloc(sizeof(float)*NumHalos);
  assert(px != NULL);
  py = (float*)malloc(sizeof(float)*NumHalos);
  assert(py != NULL);
  pz = (float*)malloc(sizeof(float)*NumHalos);
  assert(pz != NULL);
  keep = (short int*)malloc(sizeof(short int)*NumHalos);
  assert(keep != NULL);
  for(i=0;i<NumHalos;++i)
    {
      keep[i] = 1;
      px[i] = Halos[i].pos[0];
      py[i] = Halos[i].pos[1];
      pz[i] = Halos[i].pos[2];
    }
  
  searchRadius = (float) (haloFinderData.SofteningLength);
  
#ifdef LIGHTCONE
  periodic = 0;
#else
  periodic = 1;
#endif
  td = buildkdTree(px,py,pz,(int) NumHalos,NULL);
  
  //set max num nbrs
  maxNumNbrs = 100000;
  nbrsRad2 = (float*)malloc(sizeof(float)*maxNumNbrs);
  assert(nbrsRad2 != NULL);
  nbrsInd = (int*)malloc(sizeof(int)*maxNumNbrs);
  assert(nbrsInd != NULL);
  
  /*
    keep == 1 means that we keep the halo
    move down list of halos from highest mass to lowest and mark halos to get rid of by setting Parts.index = 0
  */
  for(i=0;i<NumHalos;++i)
    {
      
#ifdef DEBUG
      if(i%(NumHalos/10) == 0 && ThisTask == 0)
	fprintf(stderr,"%d of %ld (%0.2f percent) done in peak percolation\n",i,NumHalos,((double) i)/((double) NumHalos)*100.0);
#endif
      
      if(keep[i])
	{
	  NumNbrs = get_nnbrs_kdtree(Halos[i].pos,searchRadius,periodic,haloFinderData.BoxLengthSnapshot,
				     &nbrsRad2,&nbrsInd,&maxNumNbrs,px,py,pz,td);
	  
	  for(j=0;j<NumNbrs;++j)
	    if(nbrsInd[j] != i)
	      keep[nbrsInd[j]] = 0;
	}
    }
  
  //now keep all peaks not marked as junk
  haloInd = 0;
  for(i=0;i<NumHalos;++i)
    {
      if(keep[i])
	{
	  Halos[haloInd] = Halos[i];
	  ++haloInd;
	}
    }
  if(haloInd > 0)
    {
      tmpHalos = (Halo*)realloc(Halos,sizeof(Halo)*(haloInd));
      assert(tmpHalos != NULL);
      Halos = tmpHalos;
      NumHalos = haloInd;
      
      //set the halo inds
      for(i=0;i<NumHalos;++i)
	Halos[i].index = i;
    }
  else
    {
      free(Halos);
      NumHalos = 0;
      Halos = NULL;
    }
  
  free(keep);
  free(px);
  free(py);
  free(pz);
  free(nbrsRad2);
  free(nbrsInd);
  destroykdTree(td);
  
#ifdef DEBUG
  fprintf(stderr,"%d: done with peak percolation, %ld unique peaks\n",ThisTask,NumHalos);
#endif
}
